use std::collections::{HashMap, HashSet};

type Valve = &'static str;

struct Puzzle {
    total_minutes: usize,
    flow_rates: HashMap<Valve, usize>,
    shortest_paths: HashMap<(Valve, Valve), usize>,
}

fn parse_input() -> Puzzle {
    fn parse_line(value: &'static str) -> (Valve, usize, Vec<Valve>) {
        // Valve TQ has flow rate=14; tunnels lead to valves WG, KI, JE, UH, XT
        let mut iter = value.split("; ");
        let flow_rate = iter.next().unwrap();
        let tunnel_valves = iter.next().unwrap();

        let valve_name = &flow_rate[6..8];
        let flow_rate = flow_rate[23..].parse().unwrap();
        
        let tunnel_valves = 
            if tunnel_valves.starts_with("tunnels") { 
                tunnel_valves[23..]
                    .split(", ")
                    .collect()
            } else {
                tunnel_valves[22..]
                .split(", ")
                .collect()
            };

        (valve_name, flow_rate, tunnel_valves)
    }

    let mut flow_rates = HashMap::new();
    let mut tunnels = HashMap::new();

    include_str!("../input.txt")
        .lines()
        .map(parse_line)
        .for_each(|(valve_name, flow_rate, tunnel_valves)| {
            flow_rates.insert(valve_name, flow_rate);
            tunnels.insert(valve_name, tunnel_valves);
        });

    let mut shortest_paths = HashMap::new();

    for (&valve1, adjacent_valves) in tunnels.iter() {
        let mut adjacent_valves = adjacent_valves.clone();
        for distance in 1.. {
            if adjacent_valves.is_empty() {
                break;
            }

            for &valve2 in adjacent_valves.iter() {
                shortest_paths.insert((valve1, valve2), distance);
            }

            adjacent_valves = adjacent_valves
                .into_iter()
                .flat_map(|valve| tunnels.get(valve).unwrap() )
                .filter(|&valve2| !shortest_paths.contains_key(&(valve1, valve2)))
                .copied()
                .collect();
        }
    }

    Puzzle {
        total_minutes: 30,
        flow_rates,
        shortest_paths
    }

}

fn part1() {

    #[derive(Clone)]
    struct State {
        minute: usize,
        valve: Valve,
        score: usize,
        opened: HashSet<Valve>,
    }

    impl State {
        fn new() -> Self {
            Self {
                minute: 1,
                valve: "AA",
                score: 0,
                opened: HashSet::new(),
            }
        }

        fn next_states(&self, puzzle: &Puzzle) -> Vec<Self> {
            puzzle.flow_rates
                .iter()
                .filter(|(&v, &flow_rate)| (flow_rate > 0) && !self.opened.contains(v))
                .filter(|(&v, _)| {
                    let arrival_time = self.minute + puzzle.shortest_paths[&(self.valve, v)];
                    arrival_time < puzzle.total_minutes
                })
                .map(|(&v, _)| self.move_to_and_open(v, puzzle))
                .collect()
        }

        fn move_to_and_open(&self, next_valve: Valve, puzzle: &Puzzle) -> State {
            let arrival_time = self.minute + puzzle.shortest_paths.get(&(self.valve, next_valve)).unwrap();
            let mut opened = self.opened.clone();
            opened.insert(next_valve);
            State {
                minute: arrival_time + 1,
                valve: next_valve,
                score: self.score + puzzle.flow_rates[next_valve] * (puzzle.total_minutes - arrival_time),
                opened
            }
        }

        fn best_potential_score(&self, puzzle: &Puzzle) -> usize {
            self.score + 
            puzzle.flow_rates
                .iter()
                .filter(|(&v, &flow_rate)| !self.opened.contains(v) && flow_rate > 0)
                .filter_map(|(v, flow_rate)| {
                    let arrival_time = self.minute + puzzle.shortest_paths[&(self.valve, *v)];
                    (arrival_time < puzzle.total_minutes)
                        .then(|| flow_rate * (puzzle.total_minutes - arrival_time))
                })
                .sum::<usize>()
        }

    }

    fn slove(puzzle: &Puzzle, state: State, mut best_score: usize) -> usize {
        if best_score < state.best_potential_score(puzzle) {
            best_score = best_score.max(state.score);
            for next_state in state.next_states(puzzle) {
                let potentially_better_score = slove(puzzle, next_state, best_score);
                best_score = best_score.max(potentially_better_score);
            }
        }
        best_score
    }

    let puzzle = parse_input();

    println!("Part 1: {}", slove(&puzzle, State::new(), 0));
}

fn main() {
    part1();
}
